package rddl.policy.expectation;

import rddl.RDDL;
import rddl.State;
import util.Pair;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.TreeMap;

/**
 * Util class
 * Created by HuGuodong on 2017/11/24.
 */

public class Util {

    /**
     * no new Util().
     */
    private Util(){}

    /**
     * copy a MEPState mepSate from a State currentState
     * MEPState extends State
     * @param currentState
     * @return
     */
    public static MEPState coloneAndExtendState(State currentState) {

        MEPState s = new MEPState();

        s._hmPVariables = new HashMap<RDDL.PVAR_NAME, RDDL.PVARIABLE_DEF>(currentState._hmPVariables);
        s._hmTypes = new HashMap<RDDL.TYPE_NAME, RDDL.TYPE_DEF>(currentState._hmTypes);
        s._hmCPFs = new HashMap<RDDL.PVAR_NAME, RDDL.CPF_DEF>(currentState._hmCPFs);

        s._hmObject2Consts = new HashMap<RDDL.TYPE_NAME, ArrayList<RDDL.LCONST>>(currentState._hmObject2Consts);

        s._alStateNames = new ArrayList<RDDL.PVAR_NAME>(currentState._alStateNames);
        s._alActionNames = new ArrayList<RDDL.PVAR_NAME>(currentState._alActionNames);
        s._tmIntermNames = new TreeMap<Pair, RDDL.PVAR_NAME>(currentState._tmIntermNames);
        s._alIntermNames = new ArrayList<RDDL.PVAR_NAME>(currentState._alIntermNames);
        s._alObservNames = new ArrayList<RDDL.PVAR_NAME>(currentState._alObservNames);
        s._alNonFluentNames = new ArrayList<RDDL.PVAR_NAME>(currentState._alNonFluentNames);

        s._hmTypeMap = new HashMap<String, ArrayList<RDDL.PVAR_NAME>>();
        for (String key : currentState._hmTypeMap.keySet()) {
            ArrayList<RDDL.PVAR_NAME> value = currentState._hmTypeMap.get(key);
            s._hmTypeMap.put(key, new ArrayList<RDDL.PVAR_NAME>(value));
        }

        s._state = new HashMap<RDDL.PVAR_NAME, HashMap<ArrayList<RDDL.LCONST>,Object>>();
        for (RDDL.PVAR_NAME key : currentState._state.keySet()) {
            HashMap<ArrayList<RDDL.LCONST>,Object> value = currentState._state.get(key);
            s._state.put(key, new HashMap<ArrayList<RDDL.LCONST>, Object>(value));
        }

        s._nonfluents = new HashMap<RDDL.PVAR_NAME, HashMap<ArrayList<RDDL.LCONST>,Object>>();
        for (RDDL.PVAR_NAME key : currentState._nonfluents.keySet()) {
            HashMap<ArrayList<RDDL.LCONST>,Object> value = currentState._nonfluents.get(key);
            s._nonfluents.put(key, new HashMap<ArrayList<RDDL.LCONST>, Object>(value));
        }

        s._actions = new HashMap<RDDL.PVAR_NAME, HashMap<ArrayList<RDDL.LCONST>,Object>>();
        for (RDDL.PVAR_NAME key : currentState._actions.keySet()) {
            HashMap<ArrayList<RDDL.LCONST>,Object> value = currentState._actions.get(key);
            s._actions.put(key, new HashMap<ArrayList<RDDL.LCONST>, Object>(value));
        }

        s._interm = new HashMap<RDDL.PVAR_NAME, HashMap<ArrayList<RDDL.LCONST>,Object>>();
        for (RDDL.PVAR_NAME key : currentState._interm.keySet()) {
            HashMap<ArrayList<RDDL.LCONST>,Object> value = currentState._interm.get(key);
            s._interm.put(key, new HashMap<ArrayList<RDDL.LCONST>, Object>(value));
        }

        s._observ = new HashMap<RDDL.PVAR_NAME, HashMap<ArrayList<RDDL.LCONST>,Object>>();
        for (RDDL.PVAR_NAME key : currentState._observ.keySet()) {
            HashMap<ArrayList<RDDL.LCONST>,Object> value = currentState._observ.get(key);
            s._observ.put(key, new HashMap<ArrayList<RDDL.LCONST>, Object>(value));
        }

        s._alActionPreconditions = currentState._alActionPreconditions;
        s._alStateInvariants = currentState._alStateInvariants;
        s._reward = currentState._reward;
        s._nMaxNondefActions = currentState._nMaxNondefActions;

        s._nextState = new HashMap<RDDL.PVAR_NAME, HashMap<ArrayList<RDDL.LCONST>,Object>>();
        for (RDDL.PVAR_NAME key : currentState._nextState.keySet()) {
            HashMap<ArrayList<RDDL.LCONST>,Object> value = currentState._nextState.get(key);
            s._nextState.put(key, new HashMap<ArrayList<RDDL.LCONST>, Object>(value));
        }

        return s;
    }

    /**
     * find maxExpectationAction
     * @param map action,reward
     * @return
     */
    public static Map<String,Double> maxExpectationAction(Map<String,Double> map){

        Map<String,Double> rtn = new HashMap<>();

        Double maxExp = Double.MIN_VALUE;
        String action_taken = null;

        for (String action : map.keySet()) {
            Double val = map.get(action);
            if(map.get(action) > maxExp){
                maxExp = val;
                action_taken = action;
            }
        }

        rtn.put(action_taken,maxExp);

        return rtn;
    }
}
